import numpy as np
from PySide6.QtWidgets import (QMainWindow, QLabel, QScrollArea, QVBoxLayout, QWidget, QFileDialog)
from PySide6.QtGui import (QPixmap, QImage, QPainter, QPen, QColor, Qt, QAction)
from PySide6.QtCore import QPoint, Slot, Signal

from fitHalfSIN import Fit


class AnnotatableImage(QLabel):
    def __init__(self, pixmap):
        super().__init__()
        self.setPixmap(pixmap)
        self.pixmap = pixmap.copy()  # 创建原始图像的副本以进行绘制
        self.drawing = False
        self.lastPoint = QPoint()
        self.penColor = QColor(Qt.red)  # 默认绘图颜色为红色
        self.penWidth = 2  # 线条宽度
        self.current_points = []  # 用于存储当前绘制过程中所有点的坐标
        self.drawn_lines_arrays = []  # 用于存储每次绘制完成后点的坐标NumPy数组,是current_points的集合

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.drawing = True
            self.lastPoint = event.position().toPoint()
            self.current_points = [self.lastPoint]  # 开始新的绘制，重置点列表并添加起点

    def mouseMoveEvent(self, event):
        if (event.buttons() & Qt.MouseButton.LeftButton) and self.drawing:
            painter = QPainter(self.pixmap)
            pen = QPen(self.penColor, self.penWidth, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)
            painter.setPen(pen)
            current_point = event.position().toPoint()
            painter.drawLine(self.lastPoint, current_point)
            self.lastPoint = current_point
            self.setPixmap(self.pixmap)
            self.update()
            self.current_points.append(current_point)  # 添加当前点

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            self.drawing = False
            if self.current_points:  # 如果有绘制的点，则保存
                # 将点坐标列表转换为NumPy数组
                points_array = np.array([(p.x(), p.y()) for p in self.current_points])
                self.drawn_lines_arrays.append(points_array)
                # print("Drawn line points as NumPy array:\n", points_array)  # 输出本次绘制的点坐标NumPy数组
                print("曲线条数: ", len(self.drawn_lines_arrays))

                # 拟合
                self.draw_fitted_curve(Fit(points_array))
                self.update()

    def draw_fitted_curve(self, fitted_points):
        painter = QPainter(self.pixmap)
        pen = QPen(QColor(Qt.blue), self.penWidth, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)
        painter.setPen(pen)
        for i in range(len(fitted_points) - 1):
            start_point = QPoint(int(fitted_points[i][0]), int(fitted_points[i][1]))
            end_point = QPoint(int(fitted_points[i + 1][0]), int(fitted_points[i + 1][1]))
            painter.drawLine(start_point, end_point)
        # 更新 QLabel 的 pixmap
        self.setPixmap(self.pixmap)
        self.update()


class ImageWindow(QMainWindow):
    # 添加一个自定义信号，当窗口关闭时发出
    windowClosed = Signal()

    def __init__(self, image_path):
        super().__init__()

        self.pixmap = QPixmap(image_path)
        self.image_label = AnnotatableImage(self.pixmap)

        self.scroll_area = QScrollArea()
        self.scroll_area.setWidgetResizable(True)
        self.scroll_area.setWidget(self.image_label)
        self.scroll_area.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)  # 禁止横向滚动

        self.main_widget = QWidget()
        self.layout = QVBoxLayout(self.main_widget)
        self.layout.addWidget(self.scroll_area)
        self.setCentralWidget(self.main_widget)

        save_action = QAction("Save", self)
        save_action.triggered.connect(self.save_annotated_image)
        menu_bar = self.menuBar()
        file_menu = menu_bar.addMenu("File")
        file_menu.addAction(save_action)

        self.resizeEvent = self.on_resize

    @Slot()
    def on_resize(self, event=None):
        self.update_image_size()

    def update_image_size(self):
        # 获取窗口的当前宽度
        window_width = self.width()

        # 保持宽高比的情况下，将图片缩放到窗口宽度
        scaled_pixmap = self.pixmap.scaledToWidth(window_width, Qt.SmoothTransformation)

        self.image_label.pixmap = scaled_pixmap
        self.image_label.setPixmap(scaled_pixmap)
        self.image_label.adjustSize()

    def save_annotated_image(self):
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getSaveFileName(self, "Save Annotated Image", "",
                                                   "Images (*.png *.jpg);;All Files (*)", options=options)
        if file_name:
            self.image_label.pixmap.save(file_name)
            self.image_label.pixmap.save("./asset/SIN.png")

    # 重写closeEvent以发送窗口关闭信号
    def closeEvent(self, event):
        self.windowClosed.emit()  # 发送窗口关闭信号
        super().closeEvent(event)  # 调用父类的closeEvent方法


def create_image_window(image_path):
    window = ImageWindow(image_path)
    window.setWindowTitle("Annotate Image Viewer")
    return window
